#pragma once

#include <iostream>

using namespace std;

enum Colour
{
	RED,
	BLACK
};

template <class K>
struct RBNode
{
	K _data;
	RBNode<K>* _left;
	RBNode<K>* _right;
	RBNode<K>* _parent;
	Colour _col;

	RBNode(K data)
		:_data(data),
		 _left(nullptr),
		 _right(nullptr),
		 _parent(nullptr),
		 _col(RED)
	{}
};

template <class K>
class RBTree
{
public:
	using Node = RBNode<K>;

	bool Insert(K key)
	{
		Node* newnode = new Node(key);
		if (_root == nullptr)
		{
			_root = newnode;
		}
		else
		{
			Node* cur = _root;
			Node* parent = nullptr;
			while (cur)
			{
				parent = cur;
				if (cur->_data == key) return false;
				else if (cur->_data > key) cur = cur->_left;
				else cur = cur->_right;
			}

			if (parent->_data > key) parent->_left = newnode;
			else parent->_right = newnode;
			newnode->_parent = parent;

			cur = newnode;
			while (parent && parent->_col == RED)
			{
				Node* pparent = parent->_parent;
				if (parent == pparent->_left)
				{
					if (pparent->_right && pparent->_right->_col == RED)
					{
						parent->_col = BLACK;
						pparent->_right->_col = BLACK;
						pparent->_col = RED;
					}
					else
					{
						if (cur == parent->_left)
						{
							RotateR(pparent);
							pparent->_col = RED;
							parent->_col = BLACK;
						}
						else
						{
							RotateL(parent);
							RotateR(pparent);
							pparent->_col = RED;
							cur->_col = BLACK;
						}
						break;
					}
				}
				else
				{
					if (pparent->_left && pparent->_left->_col == RED)
					{
						parent->_col = BLACK;
						pparent->_left->_col = BLACK;
						pparent->_col = RED;
					}
					else
					{
						if (cur == parent->_right)
						{
							RotateL(pparent);
							pparent->_col = RED;
							parent->_col = BLACK;
						}
						else
						{
							RotateR(parent);
							RotateL(pparent);
							pparent->_col = RED;
							cur->_col = BLACK;
						}
						break;
					}
				}
				cur = pparent;
				parent = cur->_parent;
			}
		}
		_root->_col = BLACK;
		return true;
	}

	bool Check()
	{
		if (_root == nullptr) return true;
		if (_root->_col == RED) return false;

		int refNum = 0;
		Node* cur = _root;
		while (cur)
		{
			if (cur->_col == BLACK) refNum++;
			cur = cur->_left;
		}

		return Isbalance(_root, refNum, 0);
	}

	void InOrder()
	{
		_InOrder(_root);
		cout << endl;
	}

private:

	void _InOrder(Node* root)
	{
		if (root == nullptr) return;
		_InOrder(root->_left);
		cout << root->_data << " ";
		_InOrder(root->_right);
	}

	bool Isbalance(Node* root, int refNum, int blackNum)
	{
		if (root == nullptr) return refNum == blackNum;
		if (root->_col == BLACK) blackNum++;
		else
		{
			if (root->_parent && root->_parent->_col == RED)
				return false;
		}
		return (root->_left, refNum, blackNum) && (root->_right, refNum, blackNum);
	}

	void RotateR(Node* parent)
	{
		Node* subl = parent->_left;
		Node* sublr = subl->_right;
		Node* pparent = parent->_parent;
		
		if (parent == pparent->_left)
			pparent->_left = subl;
		else
			pparent->_right = subl;

		subl->_right = parent;
		subl->_parent = pparent;
		parent->_left = sublr;
		parent->_parent = subl;
		if (sublr) sublr->_parent = parent;
	}

	void RotateL(Node* parent)
	{
		Node* subr = parent->_right;
		Node* subrl = subr->_left;
		Node* pparent = parent->_parent;

		if (parent == pparent->_left)
			pparent->_left = subr;
		else
			pparent->_right = subr;

		subr->_left = parent;
		subr->_parent = pparent;
		parent->_right = subrl;
		parent->_parent = subr;
		if (subrl) subrl->_parent = parent;
	}

private:
	Node* _root = nullptr;
};